# -*- coding: UTF-8 -*-
"""
@Date    ：2025/10/23 18:30 
@Author  ：Liu Yuezhao
@Project ：bert 
@File    ：main_adapter.py
@IDE     ：PyCharm 
"""
import os
import warnings
import pandas as pd
from torch.utils.data import DataLoader
from src.tools.utils import load_config, read_json, split_data_to_train_valid_test, set_seed, undersample_negative_class
from src.transfer.time_bert_with_adapter import load_time_bert_adapter_model
from src.dataset.event_target_dataset import MBDDataset
from src.trainers.fine_tune_trainer import FineTuneTrainer

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
CONFIG_PATH = os.path.normpath(os.path.join(SCRIPT_DIR, "../../config.yaml"))
CHECKPOINT_PATH = os.path.normpath(os.path.join(SCRIPT_DIR, "../../checkpoints/pretrain/best_mlm_model.pth"))

warnings.filterwarnings("ignore")
yaml_config = load_config(CONFIG_PATH)
device = yaml_config["config"]["device"]

# 加载数据集设置
trx_dict = read_json("./data/dataset/mini_trx_unique.json")
batch_size = yaml_config["data_config"]["batch_size"]
max_seq_len = yaml_config["data_config"]["max_seq_len"]

# 设定随机种子
set_seed(yaml_config["config"]["seed"])

# 加载数据
df = pd.read_pickle("/home/datadisk/chengh/Lil/time_mbd/data/dataset/df_trx_dialogtargets_unique.pkl").sample(frac=0.1, random_state=yaml_config["config"]["seed"])

train_data, valid_data, test_data = split_data_to_train_valid_test(data=df, train_valid_ratio=[0.7, 0.2])
# 构建dataloader，这里没有做负样本下采样操作
mbd_dataset_train = MBDDataset(trx_target_data=train_data, trx_dict=trx_dict, max_seq_len=max_seq_len, device=device)
mbd_dataset_valid = MBDDataset(trx_target_data=valid_data, trx_dict=trx_dict, max_seq_len=max_seq_len, device=device)
mbd_dataset_test = MBDDataset(trx_target_data=test_data, trx_dict=trx_dict, max_seq_len=max_seq_len, device=device)

mbd_loader_train = DataLoader(mbd_dataset_train, batch_size=batch_size, shuffle=True)
mbd_loader_valid = DataLoader(mbd_dataset_valid, batch_size=batch_size, shuffle=False)
mbd_loader_test = DataLoader(mbd_dataset_test, batch_size=batch_size, shuffle=False)

# adapter模型
model = load_time_bert_adapter_model(
        checkpoint_path=CHECKPOINT_PATH,
        num_labels=2,
        adapter_name="fraud_detection",
        reduction_factor=16,
        device=device
    )

fine_tune_trainer = FineTuneTrainer(
    model=model,
    train_loader=mbd_loader_train,
    valid_loader=mbd_loader_valid,
    test_loader=mbd_loader_test,
)

fine_tune_trainer.train()
